import json
import os
from datetime import datetime, timedelta
from fastapi import APIRouter
from starlette.websockets import WebSocketDisconnect

from util.common import cut_country_id, select_file
from pathlib import Path
from routes import sqlite
import numpy as np
import pandas as pd
import win32com.client as win32
from fastapi import WebSocket


email_app = APIRouter()

# 存储所有 WebSocket 连接的列表
websockets_list = []


def classify(row,dfeu,dfyk,dfbs,dfhq2):
    classes = []
    if row['con'] in dfeu:
        classes.append('欧盟')
    if row['con'] in dfyk:
        classes.append('央库')
    if row['con'] in dfbs:
        classes.append('保税')
    if row['con'] in dfhq2:
        classes.append('元器件')
    return ','.join(classes)


# 获取要生成邮件的excel数据

@email_app.get("/getexcel")
async def getexcel():
    try:
        excelname=select_file('请选择要生成邮件的excel文件')
        if excelname=='':
            return {
                "code": 204,
                "status": "no",
                "message": "没有选择文件",
                "path": '',
                "data": []
            }
        df = pd.read_excel(excelname, sheet_name='Sheet1',
                           usecols=['*Contract No.', 'Country/Territory', 'Release Qty', 'From Warehouse',
                                    'Destination Region', 'Need Red Certificated',
                                    'DG class',
                                    'Battery Code', 'Battery Attribute', 'TD Flag(C/H)'])
        if df.empty:
            return {
                "code": 204,
                "status": "no",
                "message": "没有数据",
                "path": excelname,
                "data": []
            }
        df.rename(columns={'*Contract No.': 'con', 'Country/Territory': 'country', 'Release Qty': 'qty'}, inplace=True)
        # 按*Contract No.对Release Qty进行求和
        dfs = df.groupby(['con', 'country'])['qty'].sum().reset_index()
        # 欧盟合同
        # 在DataFrame上应用你的函数来创建一个新的列
        df['Cut Contract No'] = df['con'].apply(cut_country_id)
        # 使用条件选择来找到满足条件的行
        selected_rows = df[((df['Destination Region'] == 'European Region') & (df['Need Red Certificated'] == 'Y')) |
                           ((df['Cut Contract No'] == '826') & (df['Need Red Certificated'] == 'Y'))]

        # 将满足条件的'Contract No'放到一个列表中
        dfeu = selected_rows['con'].tolist()

        # 央库收件人
        dfyk = df[(df['From Warehouse'] == 'HQ0G1') | (df['From Warehouse'] == 'HQ0G3')][
            'con'].drop_duplicates().values
        # 保税收件人
        dfbs = df[df['From Warehouse'] == 'HQ0G2']['con'].drop_duplicates().values
        # 元器件收件人
        dfhq2 = df[df['From Warehouse'] == 'HQ2G1']['con'].drop_duplicates().values

        # 电池合同
        dft = df[(df['DG class'].notna()) | (df['Battery Code'].notna()) | (df['Battery Attribute'].notna())][
            'con'].drop_duplicates().values

        # 拍照合同
        conditions = [
            (df['TD Flag(C/H)'].str.endswith('P')),
            (df['TD Flag(C/H)'].str.endswith('Y')),
            (df['TD Flag(C/H)'].str.endswith('U'))
        ]
        values = ['拍照', 'TD', '未测算']
        df['TD Flag(C/H)'] = np.select(conditions, values, default='')
        df_ph = df[df['TD Flag(C/H)'] == '拍照']['con'].drop_duplicates().values

        dfs['reci'] = ''  # 收件人类别
        # 使用lambda函数来传递额外的参数
        dfs['reci'] = dfs.apply(lambda row: classify(row, dfeu, dfyk, dfbs, dfhq2), axis=1)

        dfs['eu'] = ''
        dfs.loc[dfs['con'].isin(dfeu), 'eu'] = '欧盟'
        dfs['bat'] = ''
        dfs.loc[dfs['con'].isin(dft), 'bat'] = '含电池'

        dfs['ph'] = ''
        dfs.loc[dfs['con'].isin(df_ph), 'ph'] = '拍照'

        return {
            "code": 200,
            "status": "ok",
            "message": "获取成功",
            "path": excelname,
            "data": dfs.to_dict(orient='records')
        }
    except Exception as e:
        return {
            "code": 500,
            "status": "error",
            "message": str(e),
            "path": excelname,
            "data": []
        }


# 生成邮件
def create_email_template(outlook, country, con, recipients_to, isph, subject, bt_html, ph_html,out_date):
    # 创建一封邮件
    mail = outlook.CreateItem(0)

    # 邮件主题
    mail.Subject = subject
    # if recipients_to==None:
    #     recipients_to = ''

    # 邮件收件人
    mail.To = recipients_to if recipients_to else ''
    head = """
<head>
<style>
    table {
        border-collapse: collapse;
    }
    th {
        background: #86A8E7;
    }
</style>
</head>
"""
    # 邮件正文（文本）
    mail.HTMLBody = f"""
    <html>
    {head}
    <body>
        <table border=1 style="border-collapse: collapse;">    
            <tr>        
                <td>紧急程度</td>        
                <td></td>    
            </tr>
            <tr>        
                <td>国家</td>        
                <td>{country}</td>    
            </tr>
            <tr>        
                <td>合同号</td>        
                <td>{con}</td>    
            </tr>
            <tr>        
                <td>CPD</td>        
                <td>{out_date}</td>    
            </tr>
            <tr>        
                <td>危险品属性</td>        
                <td></td>    
            </tr>
            <tr>        
                <td>MV物料拍照</td>        
                <td>{isph}</td>    
            </tr>
            <tr>        
                <td>运输方式</td>        
                <td></td>    
            </tr>
            <tr>        
                <td>贸易术语</td>        
                <td></td>    
            </tr>
            <tr>        
                <td>发货地址</td>        
                <td>固定</td>    
            </tr>
            <tr>        
                <td>特殊备注</td>        
                <td></td>    
            </tr>
        </table> 
        <br>
        {bt_html}
        <br>
        {ph_html}
        </body>
        </html>
    """
    # 添加附件（确保路径正确）
    attachment_path = Path(f'D:/附件/{con}.xlsx')
    mail.Attachments.Add(str(attachment_path))

    # 显示邮件前预览
    mail.Display()


async def read_excel(array_data, excelpath):
    # 获取每个国家出箱单周期数据
    current_dir = os.getcwd() # 获取当前工作目录
    countzq=pd.read_excel(os.path.join(current_dir,'出箱单周期.xlsx'),sheet_name='出箱单周期',usecols=['国家代码','出箱单周期'],dtype={'国家代码':str,'出箱单周期':int})
    if countzq.empty:
        return '出箱单周期为空'
    outlook = win32.Dispatch('Outlook.Application')
    # 读取excel
    df = pd.read_excel(excelpath, sheet_name='Sheet1',
                       dtype={'Outbound Item': str, 'Battery Code': str})
    if df.empty:
        return 'excel为空'
    # df.fillna('', inplace=True)
    df['DG class'] = df['DG class'].apply(lambda x: int(x) if isinstance(x, str) and x.isdigit() else x)
    # 复制df
    df_copy = df.copy()

    # 定义判断条件
    conditions = [
        (df_copy['TD Flag(C/H)'].str.endswith('P')),
        (df_copy['TD Flag(C/H)'].str.endswith('Y')),
        (df_copy['TD Flag(C/H)'].str.endswith('U'))
    ]
    values = ['拍照', 'TD', '未测算']
    df_copy['TD Flag'] = np.select(conditions, values, default='')

    # 电池合同
    df_bt = df[
        (df_copy['DG class'].notna()) | (df_copy['Battery Code'].notna()) | (df_copy['Battery Attribute'].notna())]

    # 拍照合同
    df_ph = df_copy[df_copy['TD Flag'] == '拍照']

    # 获取收件人
    recipients = sqlite.query("SELECT country_code,yk,bs,eu,hq2 FROM recipients")
    if recipients.empty:
        return '收件人未配置'
    recipients.rename(columns={'yk': '央库', 'bs': '保税', 'eu': '欧盟', 'hq2': '元器件'}, inplace=True)

    df = df.groupby('*Contract No.')
    for row, val in enumerate(array_data):
        for con, group in df:
            if val["con"] == con:
                group.to_excel(f'D:/附件/{con}.xlsx', index=False)
                country = array_data[row]["country"]  # 国家
                countryid = cut_country_id(con)  # 国家id
                # 计算出箱单周期：当前日期+出箱单周期
                cycle=0
                if countryid in countzq['国家代码'].values:
                    cycle = int(countzq.loc[countzq['国家代码'] == countryid, '出箱单周期'].values[0])
                # 当前日期+出箱单周期
                out_date= (datetime.now() + timedelta(days=cycle)).strftime('%Y-%m-%d')
                recipients_to = ''
                if (countryid in recipients['country_code'].values) & (val["reci"] != ''):
                    ar=val["reci"].split(',')
                    for item in ar:
                        recipients_to += str(recipients.loc[recipients['country_code'] == countryid, item].values[0])+';'
                    if recipients_to != 'None':
                        # 使用dict.fromkeys()去除重复元素并保持原始顺序
                        list_without_duplicates = list(dict.fromkeys(recipients_to.split(';')))
                        # 使用';'连接列表中的元素
                        recipients_to = ';'.join(map(str, list_without_duplicates))


                text = "";
                if array_data[row]["eu"] == '欧盟':
                    text = "欧盟RED合同，请车间注意处理";

                # 判断是否含电池

                bt_html = ''
                bttstr = array_data[row]["bat"];
                if bttstr == "含电池":
                    df_btt = df_bt[df_bt['*Contract No.'] == con][
                        ['SR No.', 'RMA Code/Ref No.', 'Outbound Item', 'Release Qty', '*Contract No.', 'DG class',
                         'Battery Code', 'Battery Attribute']]
                    # 数据库中查询Certificate Code
                    df_items = df_btt['Outbound Item'].drop_duplicates().values
                    # 将 item 列的值连接成一个字符串
                    result = ", ".join(f"'{x}'" for x in df_items)
                    sql = f"SELECT Item,CertificateCode FROM CertificateCode WHERE Item IN({result})"
                    df_res = sqlite.query(sql)
                    # 将 CertificateCode 按照 Item 进行分组，并用逗号连接
                    result = df_res.groupby('Item')['CertificateCode'].agg(','.join).reset_index()
                    df_btt = df_btt.merge(result, left_on='Outbound Item', right_on='Item', how='left')
                    df_btt.fillna('', inplace=True)
                    tr = ''
                    for index, r in df_btt.iterrows():
                        tr += f"""
                        <tr>
                        <td>{r['SR No.']}</td>
                        <td>{r['RMA Code/Ref No.']}</td>
                        <td>{r['Outbound Item']}</td>
                        <td>{r['Release Qty']}</td>
                        <td>{r['*Contract No.']}</td>
                        <td>{int(r['DG class'])}</td>
                        <td>{r['Battery Code']}</td>
                        <td>{r['Battery Attribute']}</td>
                        <td>{r['CertificateCode']}</td>
                    </tr>
                        """
                    # 电池html模板
                    bt_html = f"""
                    <h2 style="color: #007bff;">电池</h2>
                        <table  border=1>
                            <thead>
                                <tr>
                                    <th>SR No</th>
                                    <th>RMA Code/Ref No</th>
                                    <th>Outbound Item</th>
                                    <th>Release Qty</th>
                                    <th>Contract No</th>
                                    <th>DG class</th>
                                    <th>Battery Code</th>
                                    <th>Battery Attribute</th>
                                    <th>CertificateCode</th>
                                </tr>
                            </thead>
                            <tbody>
                               {tr}
                            </tbody>
                        </table>
                        """
                # 判断是否拍照
                ph_html = ''
                pzstr = array_data[row]["ph"]
                isph = 'N'  # 是否拍照
                if pzstr == "拍照":
                    isph = 'Y'
                    dfph = df_ph[df_ph['*Contract No.'] == con][
                        ['SR No.', 'RMA Code/Ref No.', 'Outbound Item', 'TD Flag', 'Release Qty', '*Contract No.',
                         'From Subinv', 'TD Flag(C/H)']]
                    df_items = dfph['Outbound Item'].drop_duplicates().values
                    # 将 item 列的值连接成一个字符串
                    result = ", ".join(f"'{x}'" for x in df_items)
                    sql = f"SELECT Item,ProductPlanner FROM SPBInformation WHERE Item IN({result})"
                    df_res = sqlite.query(sql)
                    dfph = dfph.merge(df_res, left_on='Outbound Item', right_on='Item', how='left')
                    dfph.fillna('', inplace=True)
                    tr = ''
                    for index, r in dfph.iterrows():
                        tr += f"""
                                  <tr>
                                  <td>{r['SR No.']}</td>
                                  <td>{r['RMA Code/Ref No.']}</td>
                                  <td>{r['Outbound Item']}</td>
                                   <td>{r['TD Flag']}</td>
                                  <td>{r['Release Qty']}</td>
                                  <td>{r['*Contract No.']}</td>
                                  <td>{r['From Subinv']}</td>
                                  <td>{r['TD Flag(C/H)']}</td>
                                  <td>{r['ProductPlanner']}</td>
                              </tr>
                                  """
                        ph_html = f"""
                                  <h2 style="color: #007bff;">拍照</h2>
                                      <table  border=1>
                                          <thead>
                                              <tr>
                                                  <th style="background: #D16BA5;">SR No</th>
                                                  <th style="background: #D16BA5;">RMA Code/Ref No</th>
                                                  <th style="background: #D16BA5;">Outbound Item</th>
                                                  <th style="background: #D16BA5;">TD Flag</th>
                                                  <th style="background: #D16BA5;">Release Qty</th>
                                                  <th style="background: #D16BA5;">Contract No</th>
                                                  <th style="background: #D16BA5;">From Subinv</th>
                                                  <th style="background: #D16BA5;">TD Flag(C/H)</th>
                                                  <th style="background: #D16BA5;">ProductPlanner</th>
                                              </tr>
                                          </thead>
                                          <tbody>
                                             {tr}
                                          </tbody>
                                      </table>
                                      """
                # 创建邮件模板
                constr = con + ','
                subjects = [text, bttstr, country, constr, pzstr]
                subject = (' ' * 5).join([x for x in subjects if x]) + ' '
                create_email_template(outlook, country, con, recipients_to, isph, subject, bt_html, ph_html,out_date)
                await websockets_list[0].send_text(con)
                # 跳出当前循环
                break
    return 'ok'


# pip install websockets
@email_app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    try:
        # 接受WebSocket连接
        await websocket.accept()
        websockets_list.append(websocket)
        # 接收客户端发送的文本消息
        data = await websocket.receive_text()

        # 可以将接收到的数据解析为 JSON
        data_json = json.loads(data)
        array_data = data_json["data"]
        excelpath = data_json["excelpath"]
        res = await read_excel(array_data, excelpath)
        if res == 'ok':
            await websocket.send_text('邮件创建完成')
        elif res == 'excel为空':
            await websocket.send_text('excel为空')
        elif res == '收件人未配置':
            await websocket.send_text('收件人未配置')
        elif res == '出箱单周期为空':
            await websocket.send_text('出箱单周期为空')
        else:
            await websocket.send_text('创建邮件失败')
    except WebSocketDisconnect:
        # 客户端断开连接
        websockets_list.remove(websocket)
    except Exception as e:
        # 处理其他异常
        await websocket.send_text(str(e))
    finally:
        # 关闭WebSocket连接
        await websocket.close()
        websockets_list.remove(websocket)
